package translationModel;

import java.util.AbstractMap.SimpleEntry;
import java.util.HashMap;

public class LanguageModel
{
	private HashMap<String, SimpleEntry<Double, Double>> ngrams;
	
	public LanguageModel()
	{
		ngrams = new HashMap<String, SimpleEntry<Double, Double>>();
	}
	
	public void AddNGram(String ngram, double prob, double backOffProb)
	{
		ngrams.put(ngram, new SimpleEntry<Double, Double>(prob, backOffProb));
	}
	
	public double GetScore(String word1)
	{
		return GetScore(word1, false);
	}
	
	private double GetScore(String word1, boolean backOff)
	{
		SimpleEntry<Double, Double> probs = ngrams.get(word1);
		if (probs != null)
		{
			if (!backOff)
				return probs.getKey();
			return probs.getValue();
		}
		return ngrams.get("<unk>").getKey();
	}
	
	public double GetScore(String word1, String word2)
	{
		return GetScore(word1, word2, false);
	}
	
	private double GetScore(String word1, String word2, boolean backOff)
	{
		String ngram = word1 + " " + word2;
		SimpleEntry<Double, Double> probs = ngrams.get(ngram);
		if (probs != null)
		{
			if (!backOff)
				return probs.getKey();
			return probs.getValue();
		}
		return GetScore(word2, true);
	}
	
	public double GetScore(String word1, String word2, String word3)
	{
		String ngram = word1 + " " + word2 + " " + word3;
		SimpleEntry<Double, Double> probs = ngrams.get(ngram);
		if (probs != null)
			return probs.getKey();
		return GetScore(word2, word3, true);
	}
}
